package com.shujia.mllib

import org.apache.spark.ml.classification.{NaiveBayes, NaiveBayesModel}
import org.apache.spark.ml.feature.{HashingTF, IDF, IDFModel, Tokenizer}
import org.apache.spark.sql.{DataFrame, Dataset, Row, SparkSession}

object Demo9Bayes {
  def main(args: Array[String]): Unit = {
    val spark: SparkSession = SparkSession
      .builder()
      .master("local[*]")
      .appName("Demo9Bayes")
      .config("spark.sql.shuffle.partitions", 4)
      .getOrCreate()

    import spark.implicits._
    import org.apache.spark.sql.functions._
    /**
     * 1、数据特征工程
     */
    val textDF: DataFrame = spark
      .read
      .format("csv")
      .option("sep", "\t")
      .schema("label Double,text String")
      .load("spark/data/train.txt")

    val wordsDF: DataFrame = textDF
      .as[(Double, String)]
      // 调用IK分词器 对每句话进行分词
      .map {
        case (label: Double, text: String) => {
          val wordsList: List[String] = Demo8IK.fit(text)
          (label, wordsList.mkString(" "))
        }
      }
      // 过滤掉空的数据
      .filter(_._2.stripMargin.nonEmpty)
      .toDF("label", "sentence")

    // 先调用英文的分词器 会按照空格去切分 变成TF模型需要的数据类型
    val tokenizer: Tokenizer = new Tokenizer()
      .setInputCol("sentence")
      .setOutputCol("words")
    val wordsData: DataFrame = tokenizer.transform(wordsDF)

    // 过滤完以后还剩多少句话
    //    println(wordsData.count())

    // 将包含 爱你 这个词语的句子过滤出来
    //    wordsData
    //      .select($"words")
    //      .as[Array[String]]
    //      .filter(_.contains("爱你"))
    //      .show(30000)
    // 将词语转成向量 并统计词语的出现次数
    // TF
    val hashingTF: HashingTF = new HashingTF()
      .setInputCol("words")
      .setOutputCol("rawFeatures")
      // 默认的特征数 2的18次方
      // 如果训练集很大 所有的不重复的词语的个数加起来超过了2的18次方 则可以调大
      // 一般调成2的n次方
      .setNumFeatures(262144)

    val featurizedData: DataFrame = hashingTF.transform(wordsData)

    //    featurizedData
    //      .show(200, false)

    val idf: IDF = new IDF()
      .setInputCol("rawFeatures")
      .setOutputCol("features")
    val idfModel: IDFModel = idf.fit(featurizedData)

    val rescaledData: DataFrame = idfModel.transform(featurizedData)

    //    rescaledData
    //      .show(200, false)
    /**
     * 2、将数据切分为 训练集 和 测试集
     */

    val dsArr: Array[Dataset[Row]] = rescaledData
      .randomSplit(Array(0.7, 0.3))
    val trainData: Dataset[Row] = dsArr(0)
    val testData: Dataset[Row] = dsArr(1)

    /**
     * 3、选择合适的模型
     * 文本分类 -> 贝叶斯模型
     */

    // Train a NaiveBayes model.
    val model: NaiveBayesModel = new NaiveBayes()
      .fit(trainData)

    /**
     * 4、模型评估
     */
    // Select example rows to display.
    val resDF: DataFrame = model.transform(testData)
    // 将数据中原有的label 同 预测出来的结果 prediction 进行比较 计算模型的准确率

    resDF
      .select(sum(when($"label" === $"prediction", 1).otherwise(0)) / count($"label"))
      .show()


  }

}
